"""
打板策略模块

实现各种打板策略，包括：
- 首板策略
- 连板策略
- 龙头跟随策略

作者: AI Assistant
版本: 1.0.0
日期: 2024-10-10
"""

from backtest.strategy import Strategy
import pandas as pd
import logging
from typing import Optional

logger = logging.getLogger(__name__)


class FirstBoardStrategy(Strategy):
    """首板策略 - 买入当天首次涨停的股票"""
    
    def __init__(self, 
                 min_seal_strength: float = 2.0,  # 最小封板强度(%)
                 max_limit_time: str = "10:00",    # 最晚涨停时间
                 min_turnover: float = 5.0):       # 最小换手率(%)
        """
        初始化首板策略
        
        参数:
            min_seal_strength: 最小封板强度要求(%)
            max_limit_time: 最晚涨停时间要求
            min_turnover: 最小换手率要求(%)
        """
        super().__init__("首板策略")
        self.min_seal_strength = min_seal_strength
        self.max_limit_time = max_limit_time
        self.min_turnover = min_turnover
        
        logger.info(f"初始化首板策略: 封板强度>={min_seal_strength}%, "
                   f"涨停时间<={max_limit_time}, 换手率>={min_turnover}%")
    
    def on_init(self, context):
        """策略初始化"""
        logger.info("首板策略初始化完成")
    
    def on_data(self, context, data: pd.Series):
        """
        处理行情数据
        
        参数:
            context: 策略上下文
            data: 当前K线数据
        """
        try:
            # 检查是否涨停
            if '涨跌幅' in data and data['涨跌幅'] >= 9.8:  # 接近涨停
                # 这里需要额外的涨停板数据来判断是否首板
                # 简化处理：如果涨停就买入
                
                # 检查是否已持仓
                if not context.position.holdings:
                    self.buy(context, size=0.5)  # 使用50%资金
                    logger.info(f"首板买入信号触发: {data.name}")
            
            # 持仓管理：次日卖出
            elif context.position.holdings:
                self.sell(context)
                logger.info(f"次日卖出: {data.name}")
                
        except Exception as e:
            logger.error(f"首板策略处理数据失败: {e}")


class ConsecutiveBoardStrategy(Strategy):
    """连板策略 - 买入连续涨停的股票"""
    
    def __init__(self,
                 min_consecutive: int = 2,        # 最少连板数
                 max_consecutive: int = 5,        # 最多连板数
                 min_seal_strength: float = 3.0): # 最小封板强度(%)
        """
        初始化连板策略
        
        参数:
            min_consecutive: 最少连板数要求
            max_consecutive: 最多连板数限制
            min_seal_strength: 最小封板强度要求(%)
        """
        super().__init__("连板策略")
        self.min_consecutive = min_consecutive
        self.max_consecutive = max_consecutive
        self.min_seal_strength = min_seal_strength
        
        logger.info(f"初始化连板策略: {min_consecutive}-{max_consecutive}连板, "
                   f"封板强度>={min_seal_strength}%")
    
    def on_init(self, context):
        """策略初始化"""
        self.consecutive_count = 0
        logger.info("连板策略初始化完成")
    
    def on_data(self, context, data: pd.Series):
        """
        处理行情数据
        
        参数:
            context: 策略上下文
            data: 当前K线数据
        """
        try:
            # 检查是否涨停
            if '涨跌幅' in data and data['涨跌幅'] >= 9.8:
                self.consecutive_count += 1
                
                # 判断是否满足连板条件
                if (self.min_consecutive <= self.consecutive_count <= self.max_consecutive
                    and not context.position.holdings):
                    self.buy(context, size=0.6)  # 使用60%资金
                    logger.info(f"连板买入信号触发: {self.consecutive_count}连板")
            else:
                # 没有涨停，重置计数
                self.consecutive_count = 0
                
                # 如果持仓，止损或止盈
                if context.position.holdings:
                    self.sell(context)
                    logger.info("连板中断，卖出")
                    
        except Exception as e:
            logger.error(f"连板策略处理数据失败: {e}")


class LeaderFollowStrategy(Strategy):
    """龙头跟随策略 - 跟随市场龙头股"""
    
    def __init__(self,
                 min_leader_score: float = 70.0,   # 最低龙头得分
                 hold_days: int = 3):              # 持有天数
        """
        初始化龙头跟随策略
        
        参数:
            min_leader_score: 最低龙头得分要求
            hold_days: 持有天数
        """
        super().__init__("龙头跟随策略")
        self.min_leader_score = min_leader_score
        self.hold_days = hold_days
        
        logger.info(f"初始化龙头跟随策略: 龙头得分>={min_leader_score}, "
                   f"持有{hold_days}天")
    
    def on_init(self, context):
        """策略初始化"""
        self.days_held = 0
        logger.info("龙头跟随策略初始化完成")
    
    def on_data(self, context, data: pd.Series):
        """
        处理行情数据
        
        参数:
            context: 策略上下文
            data: 当前K线数据
        """
        try:
            # 这里需要龙头股识别数据
            # 简化处理：如果涨停且成交量放大，认为是龙头
            
            if not context.position.holdings:
                # 买入条件：涨停 + 成交量放大
                if '涨跌幅' in data and data['涨跌幅'] >= 9.8:
                    if '成交量' in data and len(context.history) > 5:
                        avg_volume = context.history['成交量'].tail(5).mean()
                        if data['成交量'] > avg_volume * 1.5:
                            self.buy(context, size=0.7)  # 使用70%资金
                            self.days_held = 0
                            logger.info("龙头买入信号触发")
            else:
                # 持仓管理
                self.days_held += 1
                
                # 止损：跌破5%
                if '涨跌幅' in data and data['涨跌幅'] < -5:
                    self.sell(context)
                    logger.info(f"止损卖出: 跌幅{data['涨跌幅']:.2f}%")
                # 持有天数到达
                elif self.days_held >= self.hold_days:
                    self.sell(context)
                    logger.info(f"持有{self.days_held}天，卖出")
                    
        except Exception as e:
            logger.error(f"龙头跟随策略处理数据失败: {e}")


class AdaptiveLimitUpStrategy(Strategy):
    """自适应打板策略 - 根据市场情绪调整"""
    
    def __init__(self,
                 sentiment_threshold: int = 40):   # 市场情绪阈值
        """
        初始化自适应打板策略
        
        参数:
            sentiment_threshold: 市场情绪分数阈值
        """
        super().__init__("自适应打板策略")
        self.sentiment_threshold = sentiment_threshold
        
        logger.info(f"初始化自适应打板策略: 情绪阈值>={sentiment_threshold}")
    
    def on_init(self, context):
        """策略初始化"""
        self.market_sentiment = 50  # 默认情绪值
        logger.info("自适应打板策略初始化完成")
    
    def on_data(self, context, data: pd.Series):
        """
        处理行情数据
        
        参数:
            context: 策略上下文
            data: 当前K线数据
        """
        try:
            # 根据市场情绪决定是否参与
            if self.market_sentiment >= self.sentiment_threshold:
                # 市场情绪好，激进参与
                if not context.position.holdings:
                    if '涨跌幅' in data and data['涨跌幅'] >= 9.8:
                        self.buy(context, size=0.8)  # 使用80%资金
                        logger.info(f"自适应买入: 市场情绪{self.market_sentiment}")
            else:
                # 市场情绪差，如果持仓则退出
                if context.position.holdings:
                    self.sell(context)
                    logger.info(f"市场情绪转弱({self.market_sentiment})，卖出")
            
            # 更新市场情绪（这里简化，实际需要从市场数据计算）
            # 可以根据涨停家数、连板高度等指标计算
            
        except Exception as e:
            logger.error(f"自适应打板策略处理数据失败: {e}")

